#include "PathIntegrator.cuh"

using namespace OBR;

extern "C" __device__ Spectrum __continuation_callable__path_integrator(SamplePRD *prd)
{
    const uint mat_base = prd->params->material_base_offset;
    prd->consider_medium = false;

    CallableData *sbt_record = reinterpret_cast<CallableData *>(optixGetSbtDataPointer());
    PathIntegratorData *settings =
        reinterpret_cast<PathIntegratorData *>(sbt_record->callable_payload);

    // const params
    const uint max_depth = settings->max_depth;
    const uint rr_start = settings->rr_start;
    const float rr_max = settings->rr_max;
    const float rr_thres = settings->rr_thres;
    const float eps = 1e-4;

    uint u0, u1;
    packPointer(prd, u0, u1);
    Spectrum beta = Spectrum(1.0f);
    Spectrum L = Spectrum(0.0f);
    for (size_t depth = 0; depth < max_depth; depth++)
    {

#ifdef DEBUG_SHADER
        if (prd->log)
        {
            printf(">>>>>>>>>>>>>>>>>>> depth: %d\n", depth);
        }
#endif

        float4 urand = curand_uniform4(prd->rng_long);

        // trace current incomming ray
        optixTrace(                        // trace new ray
            prd->params->handle,           // handel
            prd->inter.wi.o,               // ray origin
            prd->inter.wi.d,               // ray direction
            eps, 1e10f,                    // ray time
            prd->scene_time,               // scene time
            OptixVisibilityMask(1),        // vis mask
            OPTIX_RAY_FLAG_DISABLE_ANYHIT, // ray flags
            0,                             // SBT offset
            1,                             // SBT stride
            0,                             // miss sbt
            u0, u1);

#ifdef DEBUG_SHADER
        if (prd->log)
        {
            printf("trace result:\n");
            prd->inter.print();
        }
#endif

        // break if beta is not valid
        if (!beta.isValid())
            break;

        // handle surface interaction
        if (!prd->hit)
        {
            /* ----------------------- case missed ---------------------- */
            L += beta * prd->bg_radiance;
            break;
        }
        else
        {
            /* -------------------- case hit surface -------------------- */
            // handle emmisive
            if (prd->inter.light_id != -1 && depth != 1)
            {
                float _pdf;
                L += beta * optixDirectCall<Spectrum, SamplePRD *, float *, int, float, float,
                                            LightCallMode>(prd->inter.light_id + 1, prd, &_pdf,
                                                           prd->inter.light_instance_id, 0, 0,
                                                           LightCallMode::CALC_LI_PDF);
                if (prd->inter.mat_id == -1)
                    break; // break when hit light without surface
            }

            // handle surface
            if (prd->inter.mat_id == -1)
            {
                /* --------------- case boundary of medium -------------- */
                prd->inter.wi.o += prd->inter.wi.d * prd->inter.wi.t;
                continue;
            }
            else
            {
                /* ----------------- case actual surface ---------------- */
                // sample direct light
                L += beta * uniformSampleOneLight(prd, urand.x);
 
                // sample material
                float pdf;
                Spectrum fr = optixDirectCall<Spectrum, SamplePRD *, float *, float, float>(
                    prd->inter.mat_id + mat_base, prd, &pdf, urand.z, urand.w,
                    MaterialCallMode::SAMPLE_FR);
                beta *= fr / pdf;

#ifdef DEBUG_SHADER
                if (prd->log)
                {
                    printf("pdf: %.6f\n", pdf);
                    printf("fr:");
                    fr.print();
                    printf("\n");
                }
#endif
            }
        }

#ifdef DEBUG_SHADER
        if (prd->log)
        {
            printf("current outgoing ray:\n");
            prd->inter.wi.print();
            printf("current beta:\n");
            beta.print();
            printf("\n");
            printf("current L:\n");
            L.print();
            printf("\n");
        }
#endif

        float max_beta_component = beta.maxComponentValue();
        if (max_beta_component > 1000.0)
            break;

        // rr
        if (depth > rr_start && max_beta_component < rr_thres)
        {
            float q = clamp(fmaxf(rr_max, 1.0f - max_beta_component), 0.0f, 1.0f);
            if (curand_uniform(prd->rng) < q)
                break;
            beta /= 1 - q;
        }
        beta = beta.clamp();
    }

    return L;
}